# 导入所需库
from sklearn.cluster import KMeans  # K-Means聚类算法
from sklearn import datasets       # 机器学习数据集
import matplotlib.pyplot as plt   # 数据可视化
import numpy as np                # 数值计算

# 加载内置的鸢尾花数据集
iris = datasets.load_iris()
X = iris.data  # 特征数据（150个样本×4个特征：花萼长宽+花瓣长宽）
y = iris.target  # 真实标签（0,1,2分别代表三种鸢尾花）

# 实例化K-Means模型，设置聚类数为3（因鸢尾花有3个种类）
km = KMeans(n_clusters=3, random_state=42)  # random_state确保结果可复现

# 训练模型并预测
km.fit(X)            # 在特征数据上训练K-Means模型
predict = km.predict(X)  # 获取每个样本的聚类标签（0,1,2）

# 创建可视化图形（10×8英寸）
plt.figure(figsize=(10, 8))

# 子图1：展示K-Means聚类结果
plt.subplot(121)  # 1行2列的第1个位置
color_map = ['orange', 'green', 'pink']  # 为3个簇分配颜色
cluster_colors = [color_map[i] for i in predict]  # 每个样本点对应的颜色
plt.title('KMeans Clustering Result')
plt.scatter(X[:, 0], X[:, 1], c=cluster_colors)  # 用花萼长宽绘制散点图
plt.xlabel('Sepal Length (cm)')  # X轴：花萼长度
plt.ylabel('Sepal Width (cm)')   # Y轴：花萼宽度

# 子图2：展示真实标签分布
plt.subplot(122)  # 1行2列的第2个位置
true_colors = [color_map[i] for i in y]  # 按真实标签着色
plt.title('True Species Distribution')
plt.scatter(X[:, 0], X[:, 1], c=true_colors)  # 相同特征维度绘图
plt.xlabel('Sepal Length (cm)')
plt.ylabel('Sepal Width (cm)')

# 优化布局并显示
plt.tight_layout()  # 自动调整子图间距
plt.show()

# 打印聚类结果标签
print("Cluster assignments:", predict)